#写一个复杂一点的模型

import torch
N,D_in,H,D_out = 64,1000,100,10

#随机创建一些训练数据
x = torch.randn(N,D_in)
y = torch.randn(N,D_out)

class TwoLayerNet(torch.nn.Module):
    def __init__(self,D_in,H,D_out):
        super(TwoLayerNet,self).__init__()
        self.linear1 = torch.nn.Linear(D_in,H)
        self.linear2 = torch.nn.Linear(H,D_out,)

    def forward(self,x):
        y_pred = self.linear2(self.linear1(x).clamp(min=0))
        return y_pred




model = TwoLayerNet(D_in,H,D_out)

model = model.to("cuda:0")

# torch.nn.init.normal_(model[0].weight)
# torch.nn.init.normal_(model[2].weight)

loss_fn = torch.nn.MSELoss(reduction = 'sum')

learning_rate = 1e-4
optimizer = torch.optim.Adam(model.parameters(),lr=learning_rate)

# learning_rate = 1e-6
# optimizer = torch.optim.SGD(model.parameters(),lr=learning_rate)


for it in range(500):
    # Forward pass
    y_pred = model(x)

    # compute loss
    loss = loss_fn(y_pred, y)
    print(it,loss.item())

    optimizer.zero_grad()

    # Backward pass
    loss.backward()

    # update weights pf w1 and w2
    optimizer.step()